#!/usr/bin/env python3
"""
Global JSON loading utilities for the data processing pipeline.

This module provides lazy loading functions for JSON configuration files
used across multiple modules in the pipeline.
"""

import json
from pathlib import Path
from typing import Dict, Any

# Global variables for lazy loading
SEGJSON = None
EMBNAMEMAP = None
GRIPPER_TRANSFORM = None
FINGER_LENGTH = None

def _load_seg_json(seg_source):
    """Lazily load the segmentation JSON file."""
    global SEGJSON
    print(f"seg_source is {seg_source} ")
    if isinstance(seg_source,bool):
        print("seg_source is bool, please use 'traverse' or 'collect' or 'eval'")
        exit()
    if seg_source:
        if SEGJSON is None:
            allowed_seg_colors_path = str(Path(__file__).parent / f'allowed_seg_colors-{seg_source}.json')
            SEGJSON = json.load(open(allowed_seg_colors_path))
        return SEGJSON
    exit()


def _load_emb_name_map():
    """Lazily load the embodiment mapping JSON file."""
    global EMBNAMEMAP
    if EMBNAMEMAP is None:
        EMBNAMEMAP = json.load(open(Path(__file__).parent / 'embodiment_mapping.json'))
    return EMBNAMEMAP

def _save_seg_json(seg_json: Dict[str, Any], seg_source):
    """Save the segmentation JSON file and update global variables."""
    global SEGJSON
    
    allowed_seg_file_name = f'allowed_seg_colors-{seg_source}.json'
    allowed_seg_colors_path = str(Path(__file__).parent / allowed_seg_file_name)
    
    with open(allowed_seg_colors_path, 'w') as f:
        json.dump(seg_json, f, indent=2)
    
    SEGJSON = seg_json

def _load_gripper_transform(embodiment: str = ''):
    """Lazily load the gripper transform JSON file."""
    global GRIPPER_TRANSFORM
    if GRIPPER_TRANSFORM is None:
        GRIPPER_TRANSFORM = json.load(open(Path(__file__).parent / f'gripper_rand_transform_{embodiment}.json'))
    return GRIPPER_TRANSFORM 

def _load_finger_length():
    """Lazily load the finger length JSON file."""
    global FINGER_LENGTH
    if FINGER_LENGTH is None:
        FINGER_LENGTH = json.load(open(Path(__file__).parent.parent / 'RoboTwin' / 'finger_length.json'))
    return FINGER_LENGTH 